#include "md.h"
#include "constants.h"
#include "parameters.h"
#include "box.h"
#include <random>
#include <fstream>
#include <iostream>
void Md::Initialize(MolSys& mol, Force& f)
{
    // clear file_to_write
    std::fstream truncfile(gMolFileWrite, std::fstream::out |std::ios_base::trunc);

    if (gIsCoordRead)
    {
        mol.ReadMolSys(gMolFileRead);
    }
    else
    {
        mol.SetAtomsInCubicBox(gAtomName, gNatoms, gBoxLength);
    }
    // if not read from file, then randomly generated
    if (!gIsVelocRead)
    {

        std::random_device rd{};
        std::mt19937 gen{rd()};
        std::normal_distribution<double> dist(0.0, 1.0);
        for (size_t i = 0; i < mol.atoms_.size(); i++)
        {
            for (size_t j = 0; j < 3; j++)
            {
                mol.atoms_[i].veloc_[j] = dist(gen);
            }
        }

        mol.RemoveCenterVeloc();

        mol.ComputeTemperature();

        // scale velocities
        for (size_t i = 0; i < mol.atoms_.size(); i++)
        {
            for (size_t j = 0; j < 3; j++)
            {
                mol.atoms_[i].veloc_[j] *= sqrt(gTemperature / mol.temperature_);
            }
        }
    }
    if (gEnsembleType == EnsembleTypes::kNve)
    {
        if (gNveMethod == NveMethods::kLeapFrog)
        {
            // move v(0) -> v(-dt/2)
            f.ComputeForce(mol);
            VelocStep(mol, -0.5 * gTimeStep);
        }
        else if (gNveMethod == NveMethods::kVelocityVerlet)
        {
            f.ComputeForce(mol);
        }
    }
    if (gEnsembleType == EnsembleTypes::kNvt)
    {
        if (gNvtMethod == NvtMethods::kLeapFrogMiddle)
        {
            // move v(0) -> v(-dt/2)
            f.ComputeForce(mol);
            VelocStep(mol, -0.5 * gTimeStep);
        }
        else if (gNvtMethod == NvtMethods::kVelocityVerletMiddle)
        {
            f.ComputeForce(mol);
        }
    }

    return;
}

void Md::Evolve(MolSys& mol, Force& f)
{
    for (int iStep = 0; iStep < gNsteps; iStep++)
    {
        if (gEnsembleType == EnsembleTypes::kNve)
        {
            NveStep(mol, f, gTimeStep);
        }
        else if (gEnsembleType == EnsembleTypes::kNvt)
        {
            NvtStep(mol, f, gTimeStep);
        }
        else if (gEnsembleType == EnsembleTypes::kNpt)
        {
            //NptStep(gTimeStep);
        }
        if (iStep % gNstepsWrite == 0)
        {
            mol.ComputeTemperature();
            mol.total_energy_ = mol.potential_energy_ + mol.kinetic_energy_;
            mol.WriteMolSys(gMolFileWrite, iStep);
        }
    }
    return;
}

void Md::VelocStep(MolSys& mol, const double &dt)
{
    for (size_t i = 0; i < mol.atoms_.size(); i++)
    {
        for (size_t j = 0; j < 3; j++)
        {
            mol.atoms_[i].veloc_[j] += mol.atoms_[i].force_[j] * dt / mol.atoms_[i].mass_;
        }
    }
    return;
}
void Md::CoordStep(MolSys& mol, const double &dt)
{
    for (size_t i = 0; i < mol.atoms_.size(); i++)
    {
        for (size_t j = 0; j < 3; j++)
        {
            mol.atoms_[i].coord_[j] += mol.atoms_[i].veloc_[j] * dt;
        }
    }
    if (gIsPeriodic)
    {
        for (int i = 0; i < mol.natoms_; i++)
        {
            PutInsideBox(mol.atoms_[i].coord_, gBoxLength);
        }
    }
    return;
}
void Md::ThermStep(MolSys& mol, const double &dt)
{
    //
    if (gBathMethod == BathMethods::kLangevin)
    {
        double c1 = exp(-gBathParameter * dt);
        double c2 = sqrt(1.0 - c1 * c1);
        std::random_device rd{};
        std::mt19937 gen{rd()};
        std::normal_distribution<double> dist(0.0, 1.0);
        for (size_t i = 0; i < mol.atoms_.size(); i++)
        {
            double v_std = sqrt(constants::Boltzmann * gTemperature / mol.atoms_[i].mass_);
            for (size_t j = 0; j < 3; j++)
            {
                mol.atoms_[i].veloc_[j] = c1 * mol.atoms_[i].veloc_[j] + c2 * v_std * dist(gen);
            }
        }
    }
    else if (gBathMethod == BathMethods::kAndersen)
    {
        std::random_device rd{};
        std::mt19937 gen{rd()};
        std::normal_distribution<double> dist_normal(0.0, 1.0);
        std::uniform_real_distribution<double> dist_uniform(0.0, 1.0);
        double val = 1.0 - exp(-gBathParameter * dt);
        for (size_t i = 0; i < mol.atoms_.size(); i++)
        {
            for (size_t j = 0; j < 3; j++)
            {
                if (dist_uniform(gen) < val)
                {
                    mol.atoms_[i].veloc_[j] = sqrt(constants::Boltzmann * gTemperature / mol.atoms_[i].mass_) * dist_normal(gen);
                }
            }
        }
    }
    else if (gBathMethod == BathMethods::kBerendsen)
    {
        mol.ComputeTemperature();
        // gBathParameter = 1/tau_t, where tau_t is usually set to 100*dt
        double scaling_factor_squared = 1.0 + dt * gBathParameter * (gTemperature / mol.temperature_ - 1.0);
        double scaling_factor = sqrt(scaling_factor_squared);
        for (size_t i = 0; i < mol.atoms_.size(); i++)
        {
            for (int j = 0; j < 3; j++)
            {
                mol.atoms_[i].veloc_[j] *= scaling_factor;
            }
        }
    }
    return;
}
// NVE for one time step
void Md::NveStep(MolSys& mol, Force& f, const double &dt)
{
    if (gNveMethod == NveMethods::kLeapFrog)
    {
        // leap-frog, v(-dt/2)->v(dt/2), x(0)->x(dt)
        f.ComputeForce(mol);
        VelocStep(mol, dt);
        CoordStep(mol, dt);
    }
    else if (gNveMethod == NveMethods::kVelocityVerlet)
    {
        VelocStep(mol, 0.5 * dt);
        CoordStep(mol, dt);
        f.ComputeForce(mol);
        VelocStep(mol, 0.5 * dt);
    }
    return;
}
//
void Md::NvtStep(MolSys& mol, Force& f, const double &dt)
{
    if (gNvtMethod == NvtMethods::kLeapFrogMiddle)
    {
        // leap-frog middle, v(-dt/2)->v(dt/2), x(0)->x(dt)
        f.ComputeForce(mol);
        VelocStep(mol, dt);
        CoordStep(mol, dt * 0.5);
        ThermStep(mol, dt);
        CoordStep(mol, dt * 0.5);
    }
    else if (gNvtMethod == NvtMethods::kVelocityVerletMiddle)
    {
        VelocStep(mol, 0.5 * dt);
        CoordStep(mol, dt * 0.5);
        ThermStep(mol, dt);
        CoordStep(mol, dt * 0.5);
        f.ComputeForce(mol);
        VelocStep(mol, 0.5 * dt);
    }
    return;
}

/*
void Md::NptStep(MolSys& mol, Force& f, const double &dt)
{
    //@TODO
    return;
}
*/
